import pandas as pd
import seaborn.objects as so
import numpy as np
import matplotlib.pyplot as plt
from seaborn import axes_style




# Download the economics dataset (from ggplot2 so comparison is apples-to-apples)
url = "https://raw.githubusercontent.com/tidyverse/ggplot2/main/data-raw/economics.csv"
economics = pd.read_csv(url)


# Quick manipulation of dataframe to convert column to datetime
df = (
    economics
    .assign(
        date = lambda df: pd.to_datetime(df['date'])
        )
)



# Default plots (Notice the xaxis only has 2 years! We'll fix this in p2)
p1 = (
    so.Plot(data=df, x='date', y='uempmed')
    .add(so.Line())
    )
p1




## Change line color and chart labels, and fix xaxis
## Note here that color is inside of the Line call, so this would color the line.
## If color were instead *inside* the so.Plot() object, SO would assign it
## a different line for each value of the factor variable (column), colored differently. (Commonly referred to as hue in seaborn)
# However, in our case, we can pass a color directly.
p2 = (
    so.Plot(data=df, x='date', y='uempmed')
    .add(so.Line(color='purple'))
    .label(title='Median Duration of Unemploymeny', x='Date', y='')
    .scale(x=so.Temporal().tick(upto=10)) #Needed for current configuration of seaborn.objects so xaxis prints more than 2 ticks
    .theme(axes_style("whitegrid")) #use a function from parent seaborn library, that will pass a prebuilt selection based on what you pass
    )

p2



## plotting multiple charts (of different line types and sizes)
p3 = (
    so.Plot(data=df)
    .add(so.Line(color='darkblue', linewidth=5), x='date', y='uempmed')
    .add(so.Line(color='red', linewidth=2, linestyle='dotted'), x='date', y='psavert')
    .label(title='Unemployment Duration (Blue)\n & Savings Rate (Red)', 
           x='Date', 
           y='')
    .scale(x=so.Temporal().tick(upto=10)) #Needed for current configuration of seaborn.objects so xaxis prints more than 2 ticks
    .theme(axes_style("whitegrid")) #use a function from parent seaborn library, that will pass a prebuilt selection based on what you pass
    )

p3


## Plotting a different line type for each group
## There isn't a natural factor in this data so let's just duplicate the data and make one up
df['fac'] = 1
df2 = df.copy()
df2['fac'] = 2
df2['uempmed'] = df2['uempmed'] - 2 + np.random.normal(size=len(df2))
df_final = pd.concat([df, df2], ignore_index=True).astype({'fac':'category'})


p4 = (
    so.Plot(data=df_final, x='date', y='uempmed', color='fac')
    .add(so.Line())
    .label(title = "Median Duration of Unemployment",
           x = "Date",
           y = "",
           color='Random Factor')
    .scale(x=so.Temporal().tick(upto=10)) #Needed for current configuration of seaborn.objects so xaxis prints more than 2 ticks
    .theme(axes_style("whitegrid")) #use a function from parent seaborn library, that will pass a prebuilt selection based on what you pass
)

p4



# Plot all 4 plots
fig, axs = plt.subplots(2, 2, figsize=(10, 8))
# Draw each plot in the corresponding subplot
p1.on(axs[0, 0]).plot()
p2.on(axs[0, 1]).plot()
p3.on(axs[1, 0]).plot()
p4.on(axs[1, 1]).plot()

# Adjust layout to avoid overlap
plt.tight_layout()

# Show the combined plot
plt.show()