"""Load, clean, and featurize AP invoice data.

Accepts either the shipped AP_INVOICE_PAYMENT.csv path or a user-uploaded
DataFrame with the same schema (see REQUIRED_COLUMNS / template.py).

Labels:
  - delay_days = Payment_Date - TERMS_DATE, defined only for PAYED invoices.
  - is_late    = delay_days > threshold (binary, for the classification objective).
Features must be knowable at invoice-approval time (no Payment_Date / Payment_Amount leakage).
"""

import numpy as np
import pandas as pd

REQUIRED_COLUMNS = [
    "CHECK_NUMBER",
    "Supplier_Number",
    "Supplier_Name",
    "Payment_Status",
    "Payment_Amount",
    "Currency",
    "Payment_Date",
    "INVOICE_NUMBER",
    "INVOICE_AMOUNT",
    "Invoice_Status",
    "INVOICE_DATE",
    "TERMS_DATE",
]

DATE_COLUMNS = ["INVOICE_DATE", "TERMS_DATE", "Payment_Date"]

FEATURE_COLS = [
    "invoice_amount_log",
    "currency",
    "supplier_freq",
    "supplier_hist_delay",
    "invoice_month",
    "invoice_day_of_week",
    "invoice_day_of_month",
    "is_month_end",
]


def validate_columns(df: pd.DataFrame):
    """Strict schema check: uploaded files must carry exactly the template's
    column names (order doesn't matter). Returns (is_valid, missing, extra)."""
    cols = set(df.columns)
    required = set(REQUIRED_COLUMNS)
    missing = sorted(required - cols)
    extra = sorted(cols - required)
    return len(missing) == 0, missing, extra


def parse_dates(df: pd.DataFrame) -> pd.DataFrame:
    """Parse date columns. Tries the shipped `DD-Mon-YY` format first (fast,
    no warnings); rows that don't match fall back to a lenient dayfirst parse
    so near-miss formats from a user-uploaded file still work."""
    df = df.copy()
    for col in DATE_COLUMNS:
        if pd.api.types.is_datetime64_any_dtype(df[col]):
            continue
        parsed = pd.to_datetime(df[col], format="%d-%b-%y", errors="coerce")
        needs_fallback = parsed.isna() & df[col].notna()
        if needs_fallback.any():
            parsed.loc[needs_fallback] = pd.to_datetime(
                df.loc[needs_fallback, col], dayfirst=True, errors="coerce"
            )
        df[col] = parsed
    return df


def load_source(source) -> pd.DataFrame:
    """`source` is either a CSV path or an already-loaded DataFrame (e.g. from
    st.file_uploader)."""
    df = source.copy() if isinstance(source, pd.DataFrame) else pd.read_csv(source)
    return parse_dates(df)


def add_label(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["delay_days"] = (df["Payment_Date"] - df["TERMS_DATE"]).dt.days
    return df


def add_classification_label(df: pd.DataFrame, threshold: float) -> pd.DataFrame:
    df = df.copy()
    df["is_late"] = (df["delay_days"] > threshold).astype(int)
    return df


def add_base_features(df: pd.DataFrame) -> pd.DataFrame:
    """Features that only need the row itself (no cross-row lookups)."""
    df = df.copy()
    df["invoice_amount_log"] = np.log1p(df["INVOICE_AMOUNT"].clip(lower=0))
    df["invoice_month"] = df["INVOICE_DATE"].dt.month
    df["invoice_day_of_week"] = df["INVOICE_DATE"].dt.dayofweek
    df["invoice_day_of_month"] = df["INVOICE_DATE"].dt.day
    df["is_month_end"] = (df["INVOICE_DATE"].dt.day >= 25).astype(int)
    df["currency"] = df["Currency"]
    return df


def build_supplier_lookup(train_df: pd.DataFrame) -> pd.DataFrame:
    """Supplier frequency + historical mean delay, fit ONLY on training rows.

    Returned lookup is applied to train/test/inference alike — this is what keeps
    the encoding leakage-safe (test-set and live invoices never contribute to it).
    """
    global_mean_delay = train_df["delay_days"].mean()
    grp = train_df.groupby("Supplier_Number")
    lookup = grp.agg(
        supplier_freq=("Supplier_Number", "count"),
        supplier_hist_delay=("delay_days", "mean"),
    ).reset_index()
    lookup.attrs["global_mean_delay"] = global_mean_delay
    lookup.attrs["global_freq"] = 1
    return lookup


def apply_supplier_lookup(df: pd.DataFrame, lookup: pd.DataFrame) -> pd.DataFrame:
    df = df.merge(lookup, on="Supplier_Number", how="left")
    df["supplier_freq"] = df["supplier_freq"].fillna(lookup.attrs.get("global_freq", 1))
    df["supplier_hist_delay"] = df["supplier_hist_delay"].fillna(
        lookup.attrs.get("global_mean_delay", df["supplier_hist_delay"].mean())
    )
    return df


def prepare_dataset(source):
    """Returns (labeled_df, unpaid_df) with base features added but WITHOUT the
    supplier lookup applied yet (that must be fit on the train split only).

    `source` is either a CSV path or an already-validated DataFrame.
    """
    df = load_source(source)
    df = add_label(df)
    df = add_base_features(df)

    labeled = df[(df["Invoice_Status"] == "APPROVED") & (df["Payment_Status"] == "PAYED")].copy()
    unpaid = df[(df["Invoice_Status"] == "APPROVED") & (df["Payment_Status"] == "NOT PAYED")].copy()

    labeled = labeled.sort_values("INVOICE_DATE").reset_index(drop=True)
    return labeled, unpaid


def time_based_split(labeled: pd.DataFrame, test_frac: float = 0.2):
    n = len(labeled)
    cut = int(n * (1 - test_frac))
    train_df = labeled.iloc[:cut].copy()
    test_df = labeled.iloc[cut:].copy()
    return train_df, test_df