import matplotlib.pyplot as plt
import pandas as pd

def evolution_2_comptes(df, id1, id2):
    def prepare_df(id):
        df_id = df[df['Registrar Account - ID'] == id].copy()
        df_id['Centralisation Date'] = pd.to_datetime(df_id['Centralisation Date'])
        df_agg = (
            df_id
            .groupby('Centralisation Date')['Quantity - AUM']
            .sum()
            .reset_index()
            .sort_values('Centralisation Date')
        )
        return df_agg

    df1 = prepare_df(id1)
    df2 = prepare_df(id2)

    plt.figure(figsize=(12, 6))

    # Courbe du premier compte
    plt.plot(df1['Centralisation Date'], df1['Quantity - AUM'],
             marker='.', linestyle='-', label=f'Account {id1}')

    # Courbe du second compte
    plt.plot(df2['Centralisation Date'], df2['Quantity - AUM'],
             marker='.', linestyle='-', label=f'Account {id2}')

    plt.title("Évolution des AUM pour deux comptes")
    plt.xlabel("Date")
    plt.ylabel("Quantity - AUM")
    plt.grid(True)
    plt.legend()           # <- important pour distinguer les comptes
    plt.tight_layout()
    plt.show()


def evolution_3_comptes(df, id1, id2, id3):
    def prepare_df(id):
        df_id = df[df['Registrar Account - ID'] == id].copy()
        df_id['Centralisation Date'] = pd.to_datetime(df_id['Centralisation Date'])
        df_agg = (
            df_id
            .groupby('Centralisation Date')['Quantity - AUM']
            .sum()
            .reset_index()
            .sort_values('Centralisation Date')
        )
        return df_agg

    df1 = prepare_df(id1)
    df2 = prepare_df(id2)
    df3 = prepare_df(id3)

    plt.figure(figsize=(12, 6))

    plt.plot(df1['Centralisation Date'], df1['Quantity - AUM'],
             marker='.', linestyle='-', label=f'Account {id1}')

    plt.plot(df2['Centralisation Date'], df2['Quantity - AUM'],
             marker='.', linestyle='-', label=f'Account {id2}')

    plt.plot(df3['Centralisation Date'], df3['Quantity - AUM'],
             marker='.', linestyle='-', label=f'Account {id3}')

    plt.title("Évolution des AUM pour trois comptes")
    plt.xlabel("Date")
    plt.ylabel("Quantity - AUM")
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.show()