"""Pre-validated universe loaders for major equity indices."""

from pathlib import Path

import polars as pl

_MASTER_FILE = Path(__file__).parent / "master_constituents.parquet"


def stoxx600() -> pl.DataFrame:
    """Load STOXX 600 universe constituents with resolved YUKKA IDs."""
    return _universe("STOXX600")


def ftse100() -> pl.DataFrame:
    """Load FTSE 100 universe constituents with resolved YUKKA IDs."""
    return _universe("FTSE100")


def nasdaq100() -> pl.DataFrame:
    """Load NASDAQ 100 universe constituents with resolved YUKKA IDs."""
    return _universe("NASDAQ100")


def sp500() -> pl.DataFrame:
    """Load S&P 500 universe constituents with resolved YUKKA IDs."""
    return _universe("SP500")


def _universe(index_name: str) -> pl.DataFrame:
    df = pl.read_parquet(_MASTER_FILE).filter(pl.col("index") == index_name)

    from_date = df.select(pl.col("start_date").min()).item()
    to_date = df.select(pl.col("end_date").max()).item()

    interval_cols = [
        col
        for col in df.columns
        if any(keyword in col.lower() for keyword in ["interval", "obser", "date", "org_permid", "index"])
    ]
    df = df.drop(interval_cols)

    # Filter out rows with null yukka_id
    excluded = df.filter(pl.col("yukka_id").is_null())
    df = df.filter(pl.col("yukka_id").is_not_null())
    if len(excluded) > 0:
        names = excluded["constituent_name"].to_list()
        print(
            f"Removed {len(excluded)} companies from the list that are not in the Yukka ontology. These include:\n"
            + "\n".join(f"  - {n}" for n in names)
        )

    print(f"List of {index_name} constituents from {from_date} to {to_date}")
    return df.unique(subset=["isin"])
