import re
import pandas as pd
import matplotlib.pyplot as plt
import plotly.graph_objects as go


def extract_year(period):
    """Converts an ONS string period (e.g., 'Nov-Jan 1993') into a chronological fractional year."""
    
    if not isinstance(period, str):
        return None

    month_offsets = {
        'Jan': 0.0, 'Feb': 1/12, 'Mar': 2/12, 'Apr': 3/12,
        'May': 4/12, 'Jun': 5/12, 'Jul': 6/12, 'Aug': 7/12,
        'Sep': 8/12, 'Oct': 9/12, 'Nov': 10/12, 'Dec': 11/12
    }

    # 1. Prevent footnotes at the bottom of the file from being parsed
    start_month = period[:3]
    
    if start_month not in month_offsets:
        return None

    # 2. Extract a valid 4-digit year
    match = re.search(r'(19\d{2}|20\d{2})', period)
    if not match:
        return None

    yr = int(match.group(1))

    # 3. Fix the calendar boundary issue for Nov/Dec
    if start_month in ['Nov', 'Dec']:
        yr -= 1

    return yr + month_offsets[start_month]


def find_col(df, *keywords):
    """Returns the first column name in df containing every keyword given."""
    for c in df.columns:
        if all(k in c for k in keywords):
            return c
    raise KeyError(f"No column found matching {keywords} in columns: {list(df.columns)}")

def parse_a05_sheet(filepath, sheet_name):
    """Parses a sheet from the ONS A05 (employment/inactivity) dataset into a clean pandas DataFrame."""
    raw = pd.read_excel(filepath, sheet_name=sheet_name, header=None)

    # Locate the start of the table
    header_start = raw.index[raw.iloc[:, 1].astype(str).str.contains("Aged 16 and over", na=False)][0]

    # Four header rows
    age_row = raw.iloc[header_start]
    measure_row = raw.iloc[header_start + 1]
    stat_row = raw.iloc[header_start + 2]
    code_row = raw.iloc[header_start + 3]

    # Data starts after the blank row
    data_start = header_start + 5
    df = raw.iloc[data_start:].reset_index(drop=True)

    # Build readable column names, forward-filling the merged age/measure header cells
    columns = []
    current_age = ""
    current_measure = ""

    for age, measure, stat, code_ in zip(age_row, measure_row, stat_row, code_row):
        if pd.notna(age):
            current_age = str(age).strip()
        if pd.notna(measure):
            current_measure = str(measure).strip()

        parts = [
            current_age,
            current_measure,
            str(stat).strip() if pd.notna(stat) else "",
            str(code_).strip() if pd.notna(code_) else "",
        ]
        columns.append("_".join(p for p in parts if p))

    df.columns = columns
    df = df.rename(columns={df.columns[0]: "Period"})
    df = df.dropna(how="all")

    # Convert numeric columns safely
    for c in df.columns[1:]:
        df[c] = df[c].astype(str).str.replace(",", "", regex=False).replace("nan", pd.NA)
        df[c] = pd.to_numeric(df[c], errors="coerce")

    df['Year_Fraction'] = df['Period'].apply(extract_year)

    # Drop rows where the period couldn't be parsed (e.g. footnotes)
    return df.dropna(subset=['Year_Fraction']).copy()


def parse_neet_sheet(filepath, sheet_name):
    """Robustly parses a sheet from the ONS NEET dataset into a clean pandas DataFrame."""
    raw = pd.read_excel(filepath, sheet_name=sheet_name, header=None)

    # Locate where the actual data starts by finding the "Levels" string in the first column
    header_start = raw.index[raw.iloc[:, 0].astype(str).str.contains("Levels", na=False)][0]

    # Headers are spread across three rows directly above the data start
    age_row = raw.iloc[header_start - 2]
    metric_row = raw.iloc[header_start - 1]
    sub_metric_row = raw.iloc[header_start]

    data_start = header_start + 3
    df = raw.iloc[data_start:].reset_index(drop=True)

    columns = []
    current_age = ""
    current_metric = ""

    # Forward-fill logic for the merged header cells
    for age, metric, sub in zip(age_row, metric_row, sub_metric_row):
        if pd.notna(age) and str(age).strip() != "":
            current_age = str(age).strip()
        if pd.notna(metric) and str(metric).strip() != "":
            current_metric = str(metric).strip()

        parts = [current_age, current_metric, str(sub).strip() if pd.notna(sub) else ""]
        columns.append("_".join(p for p in parts if p and p != 'NaN'))

    df.columns = columns
    df = df.rename(columns={df.columns[0]: "Period"})
    df = df.dropna(how="all", subset=["Period"])

    # Clean numerical columns and handle '..' missing-data markers
    for c in df.columns[1:]:
        df[c] = df[c].astype(str).str.replace(",", "", regex=False).replace("..", pd.NA).replace("nan", pd.NA)
        df[c] = pd.to_numeric(df[c], errors="coerce")

    df['Year_Fraction'] = df['Period'].apply(extract_year)
    df = df.dropna(subset=['Year_Fraction']).copy()

    # The NEET dataset begins tracking later than A05, and the early rows for the
    # 18-24 rate are blank; drop them so charts start from real data.
    rate_col = find_col(df, "18-24", "NEET as a percentage")
    return df.dropna(subset=[rate_col]).copy()

# Load datasets

# A05 workbook: employment & inactivity rates
df_people = parse_a05_sheet("data/A05.xls", "People")
df_men = parse_a05_sheet("data/A05.xls", "Men")
df_women = parse_a05_sheet("data/A05.xls", "Women")

# NEET workbook: seasonally-adjusted NEET rates
df_neet_people = parse_neet_sheet("data/NEET.xlsx", "People - SA")
df_neet_men = parse_neet_sheet("data/NEET.xlsx", "Men - SA")
df_neet_women = parse_neet_sheet("data/NEET.xlsx", "Women - SA")

# Here you can see all columns we have available to us

dfs = [df_people, df_men, df_women, df_neet_people, df_neet_men, df_neet_women]

for df in dfs:
    for column in df.columns:
        print(column)
    print("\n\n")

col_men_emp = find_col(df_men, "Aged 18-24", "Employment_rate")
col_men_inac = find_col(df_men, "Aged 18-24", "Inactivity_rate")
col_women_emp = find_col(df_women, "Aged 18-24", "Employment_rate")
col_women_inac = find_col(df_women, "Aged 18-24", "Inactivity_rate")

# Employment: Men vs Women
fig_emp = go.Figure(data=[
    line_trace(df_men['Year_Fraction'], df_men[col_men_emp], 'Men (18-24)', '#1f77b4'),
    line_trace(df_women['Year_Fraction'], df_women[col_women_emp], 'Women (18-24)', '#e377c2'),
])
style_figure(fig_emp, 'Employment Rates for 18-24 Year Olds: Men vs Women',
             'Employment Rate (%)', yaxis_range=[0, 100], width=800)
fig_emp.show()

# Inactivity: Men vs Women
fig_inac = go.Figure(data=[
    line_trace(df_men['Year_Fraction'], df_men[col_men_inac], 'Men (18-24)', '#1f77b4'),
    line_trace(df_women['Year_Fraction'], df_women[col_women_inac], 'Women (18-24)', '#e377c2'),
])
style_figure(fig_inac, 'Inactivity Rates for 18-24 Year Olds: Men vs Women',
             'Inactivity Rate (%)', yaxis_range=[0, 100], width=800)
fig_inac.show()

col_neet_1824 = find_col(df_neet_people, "18-24", "NEET as a percentage")

fig_neet_overall = go.Figure(data=[
    line_trace(df_neet_people['Year_Fraction'], df_neet_people[col_neet_1824],
               'Aged 18-24 NEET Rate', '#d62728', width=2.5)
])
style_figure(fig_neet_overall, 'Overall NEET Rates in the UK (Seasonally Adjusted)', 'NEET Rate (%)')
fig_neet_overall.show()

fig_neet_gender = go.Figure(data=[
    line_trace(df_neet_men['Year_Fraction'], df_neet_men[col_neet_1824], 'Men (18-24) NEET Rate', '#1f77b4'),
    line_trace(df_neet_women['Year_Fraction'], df_neet_women[col_neet_1824], 'Women (18-24) NEET Rate', '#e377c2'),
])
style_figure(fig_neet_gender, 'NEET Rates for 18-24 Year Olds: Men vs Women', 'NEET Rate (%)')
fig_neet_gender.show()

col_emp_1824 = find_col(df_people, "Aged 18-24", "Employment_rate")

df_analysis = pd.merge(df_people, df_neet_people, on="Year_Fraction", suffixes=('_A05', '_NEET'))
df_analysis = df_analysis[['Year_Fraction', 'Period_A05', col_emp_1824, col_neet_1824]].dropna().copy()

df_analysis['Potential_Emp_Half_NEET'] = df_analysis[col_emp_1824] + (df_analysis[col_neet_1824] * 0.5)
df_analysis['Potential_Emp_Zero_NEET'] = df_analysis[col_emp_1824] + df_analysis[col_neet_1824]

fig_scenario = go.Figure(data=[
    line_trace(df_analysis['Year_Fraction'], df_analysis[col_emp_1824],
               'Actual Employment Rate', '#1f77b4', width=2.5),
    line_trace(df_analysis['Year_Fraction'], df_analysis['Potential_Emp_Half_NEET'],
               'If NEET Rate Halved', '#ff7f0e', dash='dot'),
    go.Scatter(
        x=df_analysis['Year_Fraction'], y=df_analysis['Potential_Emp_Zero_NEET'],
        mode='lines', name='If NEET Rate was 0% (All NEETs Employed)',
        line=dict(color='#2ca02c', width=2.5),
        fill='tonexty',  # shades the area down to the previous trace
        fillcolor='rgba(44, 160, 44, 0.1)',
    ),
])
style_figure(
    fig_scenario, 'The "NEET Gap": Impact of NEET Rates on 18-24 Employment', 'Employment Rate (%)',
    height=600, width=900,
    legend=dict(yanchor="top", y=0.99, xanchor="left", x=0.01, bgcolor="rgba(255, 255, 255, 0.8)"),
)
fig_scenario.show()

target_col = find_col(df_people, "Aged 18-24", "Inactivity_rate")

plt.figure(figsize=(10, 6))
plt.plot(df_people['Year_Fraction'], df_people[target_col], color='#d62728', linewidth=2.5)
plt.title('Economic Inactivity Rate: 18-24 Year Olds (UK)', fontsize=14, pad=15)
plt.xlabel('Year', fontsize=12)
plt.ylabel('Inactivity Rate (%)', fontsize=12)
plt.grid(True, alpha=0.3)
plt.tight_layout()

plt.savefig('inactivity_1824.png')
plt.show()

latest_emp = df_people.iloc[-1][col_emp_1824]

pop_col = find_col(df_neet_people, "18-24", "Total people in relevant population group")
unemp_col = find_col(df_neet_people, "18-24", "Unemployed")
inact_col = find_col(df_neet_people, "18-24", "Economically inactive")

df_neet_people['Unemployed_NEET_Rate'] = (df_neet_people[unemp_col] / df_neet_people[pop_col]) * 100
df_neet_people['Inactive_NEET_Rate'] = (df_neet_people[inact_col] / df_neet_people[pop_col]) * 100

latest_neet = df_neet_people.iloc[-1]
unemp_neet = latest_neet['Unemployed_NEET_Rate']
inact_neet = latest_neet['Inactive_NEET_Rate']

# The remaining share is assumed to be in full-time education (inactive, non-NEET)
inact_students = 100 - latest_emp - (unemp_neet + inact_neet)

labels = ['Employed', 'Full-Time Education (Inactive)', 'NEET (Economically Inactive)', 'NEET (Unemployed)']
values = [latest_emp, inact_students, inact_neet, unemp_neet]
colors = ['#2ca02c', '#1f77b4', '#9467bd', '#d62728']

fig_pie = go.Figure(data=[go.Pie(
    labels=labels,
    values=values,
    pull=[0, 0, 0.1, 0.1],  # explode the NEET slices slightly for emphasis
    marker=dict(colors=colors, line=dict(color='#000000', width=1)),
    textinfo='percent+label',
    textposition='outside',
)])

fig_pie.update_layout(
    title_text='UK Youth (18-24) Economic Status Breakdown (Latest Data)',
    height=600,
    showlegend=False,
)
fig_pie.show()