import pandas as pd
import numpy as np

from pathlib import Path
module_dir = Path(__file__).parent

def add_possible_swimups(DF):
    ec = DF.EventClass.values[0]
    if ec in [2,3,4,5,6,7,8,9,10,12,13]:
        sc = [ec,ec-1]
    else:
        sc = [ec]
    return DF.assign(SportClass=[sc]*len(DF))

def set_swimups(DF):
    inclass = DF[DF.EventClass == DF.SportClass]
    swimup = DF[~DF.index.isin(inclass) & ~DF.SportClass.isin(DF.EventClass)]
    return pd.concat([inclass,swimup])


WC2025_MQS_MET = (
    pd.read_excel(module_dir / '2025 WPS WC MQS and METs.xlsx')
)
WC2025_Events = WC2025_MQS_MET.drop('MQS MET'.split(),axis=1)

LA2028_Events = (
    pd.read_excel(module_dir / 'Draft LA2028 Events.xlsx')
    .ffill()
    .groupby('EventClass',group_keys=False)
    .apply(add_possible_swimups)
    .explode('SportClass')
    .reset_index(drop=True)
    .groupby('Gender Distance Stroke'.split(),group_keys=False)
    .apply(set_swimups)
    .sort_values('Gender Stroke Distance EventClass SportClass'.split())
    .reset_index(drop=True)
    .astype({'Distance':int,'EventClass':int,'SportClass':int})
)

__all__ = 'WC2025_MQS_MET WC2025_Events LA2028_Events'.split()
