"""Panel data container and manipulation utilities.

Provides the EAPPanel class for structured panel operations, standalone
functions for lagging and aligning characteristics with returns, and
load_gfd_chars() for pulling firm characteristics from the research PostgreSQL
database (wrds.wrds_global_factor, Global Factor Data / Chen-Zimmerman format).
"""

from __future__ import annotations

from collections.abc import Callable
from typing import Any, Literal

import numpy as np
import pandas as pd
import statsmodels.api as sm


class EAPPanel:
    """Panel data container with MultiIndex (date, entity).

    Wraps a long-format DataFrame and provides vectorized panel operations
    including shifting, forward-filling, rolling statistics, and rolling
    regressions. All within-entity operations prevent cross-entity leakage.

    Inspired by pyanomaly.Panel. Reimplemented with full type annotations
    and no external dependency on pyanomaly.

    Attributes
    ----------
    data : pd.DataFrame
        MultiIndex (date_col, entity_col) DataFrame.
    freq : str
        'M' (monthly) or 'D' (daily).
    date_col : str
        Name of the date level in the MultiIndex.
    entity_col : str
        Name of the entity level in the MultiIndex.
    """

    def __init__(
        self,
        data: pd.DataFrame,
        date_col: str = "date",
        entity_col: str = "permno",
        freq: Literal["M", "D"] = "M",
    ) -> None:
        """Initialize EAPPanel from a DataFrame.

        Parameters
        ----------
        data : pd.DataFrame
            If MultiIndex with (date_col, entity_col): used directly.
            If flat columns: sets the MultiIndex from date_col and entity_col.
        date_col : str
            Date level name.
        entity_col : str
            Entity level name.
        freq : str
            Panel frequency: 'M' (monthly) or 'D' (daily).
        """
        self.date_col = date_col
        self.entity_col = entity_col
        self.freq = freq

        # If data already has the correct MultiIndex, use it directly
        if (
            isinstance(data.index, pd.MultiIndex)
            and data.index.names == [date_col, entity_col]
        ):
            self.data = data.copy()
        else:
            # Set MultiIndex from columns
            df = data.copy()
            self.data = df.set_index([date_col, entity_col])

        # Sort the index for consistent groupby operations
        self.data = self.data.sort_index()

    # ---------------------------------------------------------------
    # Construction
    # ---------------------------------------------------------------

    @classmethod
    def from_long(
        cls,
        data: pd.DataFrame,
        date_col: str = "date",
        entity_col: str = "permno",
    ) -> EAPPanel:
        """Construct an EAPPanel from a long-format DataFrame.

        Parameters
        ----------
        data : pd.DataFrame
            Long-format panel with date and entity columns.
        date_col : str
            Name of the date column.
        entity_col : str
            Name of the entity column.

        Returns
        -------
        EAPPanel
        """
        return cls(data=data, date_col=date_col, entity_col=entity_col)

    # ---------------------------------------------------------------
    # Panel operations
    # ---------------------------------------------------------------

    def shift(self, cols: list[str], n: int = 1) -> EAPPanel:
        """Lag columns by n periods within each entity.

        Shifting is done within-entity using groupby to prevent
        cross-entity leakage. The first n periods for each entity
        become NaN (no forward information leakage).

        Parameters
        ----------
        cols : list[str]
            Columns to shift.
        n : int
            Number of periods to lag (positive = backward shift).

        Returns
        -------
        EAPPanel
            New panel with shifted columns.
        """
        result = self.data.copy()

        # Group by entity level and shift within each group
        for col in cols:
            result[col] = result.groupby(level=self.entity_col)[col].shift(n)

        return EAPPanel(
            data=result,
            date_col=self.date_col,
            entity_col=self.entity_col,
            freq=self.freq,
        )

    def populate(
        self,
        method: Literal["ffill", "bfill"] = "ffill",
        limit: int | None = None,
    ) -> EAPPanel:
        """Forward-fill (or back-fill) missing values within each entity.

        Operates across the full date grid per entity, filling gaps
        in the time series without crossing entity boundaries.

        Parameters
        ----------
        method : str
            Fill method: 'ffill' (forward) or 'bfill' (backward).
        limit : int | None
            Maximum number of consecutive periods to fill.

        Returns
        -------
        EAPPanel
            New panel with filled values.
        """
        # Apply fill method within each entity group
        if method == "ffill":
            filled = self.data.groupby(level=self.entity_col).ffill(limit=limit)
        else:
            filled = self.data.groupby(level=self.entity_col).bfill(limit=limit)

        return EAPPanel(
            data=filled,
            date_col=self.date_col,
            entity_col=self.entity_col,
            freq=self.freq,
        )

    def rolling(
        self,
        cols: list[str],
        n: int,
        func: Callable[..., Any],
        min_n: int | None = None,
    ) -> pd.DataFrame:
        """Apply a rolling window function within each entity.

        Parameters
        ----------
        cols : list[str]
            Columns to compute the rolling statistic on.
        n : int
            Window size in periods.
        func : Callable
            Function to apply to each rolling window (e.g. np.mean).
        min_n : int | None
            Minimum number of non-NaN observations required.
            Defaults to n if not specified.

        Returns
        -------
        pd.DataFrame
            DataFrame with the same index, containing rolling results.
        """
        min_periods = min_n if min_n is not None else n
        result_frames: list[pd.DataFrame] = []

        for col in cols:
            # Rolling within each entity to prevent cross-entity leakage
            rolled = (
                self.data[col]
                .groupby(level=self.entity_col)
                .rolling(window=n, min_periods=min_periods)
                .apply(func, raw=True)
            )
            # The groupby + rolling produces an extra entity level; drop it
            rolled = rolled.droplevel(0)
            result_frames.append(rolled.rename(col).to_frame())

        return pd.concat(result_frames, axis=1)

    def rolling_regression(
        self,
        y_col: str,
        x_cols: list[str],
        n: int,
        min_n: int | None = None,
        add_const: bool = True,
    ) -> pd.DataFrame:
        """Rolling OLS regression per entity.

        Estimates a linear regression in a rolling window for each entity
        independently, preventing cross-entity information leakage.

        Parameters
        ----------
        y_col : str
            Dependent variable column.
        x_cols : list[str]
            Independent variable columns.
        n : int
            Rolling window size.
        min_n : int | None
            Minimum observations for a valid regression. Defaults to n.
        add_const : bool
            If True, adds a constant (intercept) to the regression.

        Returns
        -------
        pd.DataFrame
            Coefficients for each observation. Columns are 'const' (if
            add_const) plus x_cols.
        """
        min_periods = min_n if min_n is not None else n

        # Determine output column names
        coef_names = (["const"] if add_const else []) + list(x_cols)
        # Pre-allocate result with NaN
        coef_df = pd.DataFrame(
            np.nan,
            index=self.data.index,
            columns=coef_names,
        )

        # Run rolling OLS within each entity
        for entity, group in self.data.groupby(level=self.entity_col):
            y_full = group[y_col]
            x_full = group[x_cols]

            # Iterate over rolling windows
            for i in range(len(group)):
                if i + 1 < min_periods:
                    continue

                start = max(0, i + 1 - n)
                y_win = y_full.iloc[start : i + 1]
                x_win = x_full.iloc[start : i + 1]

                # Drop rows where y or any x is NaN
                valid = y_win.notna() & x_win.notna().all(axis=1)
                y_clean = y_win[valid]
                x_clean = x_win[valid]

                if len(y_clean) < min_periods:
                    continue

                # Add constant if requested
                x_reg = sm.add_constant(x_clean) if add_const else x_clean

                try:
                    model = sm.OLS(y_clean, x_reg).fit()
                    coef_df.loc[group.index[i], coef_names] = model.params.values
                except Exception:
                    # Singular matrix or other numerical issues: leave as NaN
                    pass

        return coef_df

    def apply_to_dates(
        self,
        func: Callable[[pd.DataFrame], pd.DataFrame],
    ) -> pd.DataFrame:
        """Apply a cross-sectional function to each date group.

        Parameters
        ----------
        func : Callable
            Function that takes a DataFrame (entities as rows) and returns
            a DataFrame with the same index.

        Returns
        -------
        pd.DataFrame
            Concatenated results across all dates.
        """
        results: list[pd.DataFrame] = []

        for _date, group in self.data.groupby(level=self.date_col):
            results.append(func(group))

        return pd.concat(results)

    def apply_to_entities(
        self,
        func: Callable[[pd.DataFrame], pd.DataFrame],
    ) -> pd.DataFrame:
        """Apply a time-series function to each entity group.

        Parameters
        ----------
        func : Callable
            Function that takes a DataFrame (dates as rows) and returns
            a DataFrame with the same index.

        Returns
        -------
        pd.DataFrame
            Concatenated results across all entities.
        """
        results: list[pd.DataFrame] = []

        for _entity, group in self.data.groupby(level=self.entity_col):
            results.append(func(group))

        return pd.concat(results)

    # ---------------------------------------------------------------
    # Utilities
    # ---------------------------------------------------------------

    def to_long(self) -> pd.DataFrame:
        """Convert the panel back to a flat long-format DataFrame.

        Returns
        -------
        pd.DataFrame
            DataFrame with date_col and entity_col as regular columns.
        """
        return self.data.reset_index()

    def merge(
        self,
        right: EAPPanel | pd.DataFrame,
        on: list[str] | None = None,
        how: Literal["left", "right", "inner", "outer"] = "left",
    ) -> EAPPanel:
        """Merge another panel or DataFrame into this one.

        Parameters
        ----------
        right : EAPPanel | pd.DataFrame
            Data to merge. If EAPPanel, its underlying data is used.
        on : list[str] | None
            Columns to join on. If None, joins on the MultiIndex.
        how : str
            Join type: 'left', 'right', 'inner', 'outer'.

        Returns
        -------
        EAPPanel
            New panel with merged data.
        """
        # Extract the DataFrame from an EAPPanel if needed
        right_df = right.data if isinstance(right, EAPPanel) else right

        if on is not None:
            # Reset index to merge on regular columns
            left_flat = self.data.reset_index()
            right_flat = (
                right_df.reset_index()
                if isinstance(right_df.index, pd.MultiIndex)
                else right_df
            )
            merged = left_flat.merge(right_flat, on=on, how=how)
            return EAPPanel(
                data=merged,
                date_col=self.date_col,
                entity_col=self.entity_col,
                freq=self.freq,
            )
        else:
            # Join on the MultiIndex directly
            merged = self.data.join(right_df, how=how)
            return EAPPanel(
                data=merged,
                date_col=self.date_col,
                entity_col=self.entity_col,
                freq=self.freq,
            )

    def inspect(self) -> dict[str, Any]:
        """Return summary statistics about the panel.

        Returns
        -------
        dict[str, Any]
            Keys: n_entities, n_dates, n_obs, pct_missing, unbalanced.
        """
        # Count unique entities and dates from the MultiIndex
        n_entities = self.data.index.get_level_values(self.entity_col).nunique()
        n_dates = self.data.index.get_level_values(self.date_col).nunique()
        n_obs = len(self.data)

        full_grid = n_entities * n_dates
        pct_missing = 1.0 - (n_obs / full_grid) if full_grid > 0 else 0.0

        # Check balance: each entity should have the same number of dates
        obs_per_entity = self.data.groupby(level=self.entity_col).size()
        unbalanced = int(obs_per_entity.nunique() > 1)

        return {
            "n_entities": n_entities,
            "n_dates": n_dates,
            "n_obs": n_obs,
            "pct_missing": round(pct_missing, 6),
            "unbalanced": unbalanced,
        }


# -------------------------------------------------------------------
# Standalone functions
# -------------------------------------------------------------------


def lag_panel(
    data: pd.DataFrame,
    cols: list[str],
    n_periods: int = 1,
    entity_col: str = "permno",
    date_col: str = "date",
) -> pd.DataFrame:
    """Lag columns by n_periods within each entity.

    Prevents cross-entity leakage by grouping on entity_col before
    shifting. The first n_periods observations per entity become NaN.

    Parameters
    ----------
    data : pd.DataFrame
        Long-format panel DataFrame.
    cols : list[str]
        Columns to lag.
    n_periods : int
        Number of periods to shift (positive = backward).
    entity_col : str
        Entity identifier column.
    date_col : str
        Date column (used for sorting within entity).

    Returns
    -------
    pd.DataFrame
        Copy of input with lagged columns.
    """
    result = data.copy()

    # Sort by entity and date to ensure correct temporal ordering
    result = result.sort_values([entity_col, date_col])

    for col in cols:
        result[col] = result.groupby(entity_col)[col].shift(n_periods)

    return result


def load_gfd_chars(
    chars: list[str] | None = None,
    start_date: str | None = None,
    end_date: str | None = None,
    countries: list[str] | None = None,
    include_meta: bool = True,
    db_name: str = "research",
    schema: str = "wrds",
    table: str = "wrds_global_factor",
) -> pd.DataFrame:
    """Load firm characteristics from the Global Factor Data table.

    Queries ``{schema}.{table}`` in the ``{db_name}`` PostgreSQL database
    (default: research.wrds.wrds_global_factor, Chen-Zimmerman convention).

    Column names follow the GFD/Chen-Zimmerman convention, which differs from
    the JKP/OSAP names in JKP_153. Use ``eapctf.utils.constants.GFD_CHAR_MAP``
    to map between naming conventions.

    Always-included columns
    -----------------------
    Identifiers : ``iid`` (firm ID), ``eom`` (date, Period[M]), ``excntry``,
        ``permno``, ``gvkey``
    Targets     : ``ret``, ``ret_exc``, ``ret_exc_lead1m``

    When ``include_meta=True`` (default):
        ``me``, ``me_company``,
        ``sic``, ``naics``, ``ff49``, ``gics``,
        ``obs_main``, ``primary_sec``, ``common``, ``size_grp``

    Parameters
    ----------
    chars : list[str] | None
        Additional GFD column names to pull (e.g. the 153 characteristics).
        Use GFD column names (not JKP names). If None, only the always-included
        columns above are returned.
    start_date : str | None
        Start filter as 'YYYY-MM-DD' (inclusive), applied to the ``eom`` column.
    end_date : str | None
        End filter as 'YYYY-MM-DD' (inclusive).
    countries : list[str] | None
        ISO 3166-1 alpha-3 codes (e.g. ``["USA"]``). None = all countries.
    include_meta : bool, default True
        Whether to include market-equity, industry classification, and
        sample-selection flag columns listed above.
    db_name : str
        PostgreSQL database name. Default: "research".
    schema : str
        Schema name. Default: "wrds".
    table : str
        Table name. Default: "wrds_global_factor".

    Returns
    -------
    pd.DataFrame
        Long-format panel sorted by ``eom``, ``iid``. The ``eom`` column is
        converted to ``pd.Period[M]``.

    Raises
    ------
    ImportError
        If psycopg2 is not installed.
    RuntimeError
        If the database connection fails.

    Notes
    -----
    The table is ~29 GB (1925-2025, 444 columns). Always specify ``chars``
    and date/country filters to avoid loading the entire table.

    Examples
    --------
    Load US JKP characteristics from 2000 onward:

    >>> from eapctf.utils.constants import GFD_CHAR_MAP, JKP_153, GFD_CHAR_UNAVAILABLE
    >>> avail = [c for c in JKP_153 if c not in GFD_CHAR_UNAVAILABLE]
    >>> gfd_cols = list({GFD_CHAR_MAP.get(c, c) for c in avail})
    >>> df = load_gfd_chars(chars=gfd_cols, start_date="2000-01-01", countries=["USA"])
    """
    try:
        import psycopg2
    except ImportError as exc:
        raise ImportError(
            "psycopg2 is required for load_gfd_chars(). "
            "Install with: pip install psycopg2-binary"
        ) from exc

    # Always-included: identifiers + targets
    id_cols = ["iid", "eom", "excntry", "permno", "gvkey"]
    target_cols = ["ret", "ret_exc", "ret_exc_lead1m"]

    # Optional meta columns: market equity, industry classification, filter flags
    meta_cols: list[str] = []
    if include_meta:
        meta_cols = [
            "me", "me_company",
            "sic", "naics", "ff49", "gics",
            "obs_main", "primary_sec", "common", "size_grp",
        ]

    # User-requested characteristic columns
    char_cols: list[str] = list(chars) if chars is not None else []

    # Deduplicate while preserving logical order
    seen: set[str] = set()
    select_cols: list[str] = []
    for col in id_cols + target_cols + meta_cols + char_cols:
        if col not in seen:
            seen.add(col)
            select_cols.append(col)

    col_sql = ", ".join(f'"{c}"' for c in select_cols)

    # Build WHERE clause
    conditions: list[str] = []
    params: list[Any] = []

    if start_date is not None:
        conditions.append('"eom" >= %s::date')
        params.append(start_date)
    if end_date is not None:
        conditions.append('"eom" <= %s::date')
        params.append(end_date)
    if countries is not None:
        placeholders = ", ".join(["%s"] * len(countries))
        conditions.append(f'"excntry" IN ({placeholders})')
        params.extend(countries)

    where_clause = f"WHERE {' AND '.join(conditions)}" if conditions else ""
    query = (
        f'SELECT {col_sql} FROM "{schema}"."{table}" {where_clause}'
        f' ORDER BY "eom", "iid"'
    )

    try:
        conn = psycopg2.connect(dbname=db_name)
    except psycopg2.OperationalError as exc:
        raise RuntimeError(
            f"Cannot connect to database '{db_name}'. "
            "Ensure PostgreSQL is running and accessible."
        ) from exc

    try:
        df = pd.read_sql(query, conn, params=params if params else None)
    finally:
        conn.close()

    # Convert eom to Period[M] for month-level panel operations
    if "eom" in df.columns:
        df["eom"] = pd.to_datetime(df["eom"]).dt.to_period("M")

    return df


def align_returns_chars(
    chars: pd.DataFrame,
    returns: pd.DataFrame,
    lag: int = 1,
    entity_col: str = "permno",
    date_col: str = "date",
) -> pd.DataFrame:
    """Merge lagged characteristics with forward returns.

    Characteristics are lagged by ``lag`` periods so that at each date t,
    the characteristics from t-lag are paired with the return at t.
    This prevents look-ahead bias in predictive regressions.

    Parameters
    ----------
    chars : pd.DataFrame
        Characteristics panel with entity and date columns.
    returns : pd.DataFrame
        Returns panel with entity, date, and return columns.
    lag : int
        Number of periods to lag characteristics (default 1).
    entity_col : str
        Entity identifier column.
    date_col : str
        Date column.

    Returns
    -------
    pd.DataFrame
        Merged panel where each row has return at t and chars from t-lag.
    """
    # Identify characteristic columns (everything except entity and date)
    char_cols = [c for c in chars.columns if c not in [entity_col, date_col]]

    # Lag characteristics by the specified number of periods
    chars_lagged = lag_panel(
        chars, cols=char_cols, n_periods=lag,
        entity_col=entity_col, date_col=date_col,
    )

    # Merge lagged characteristics with returns on entity and date
    merged = returns.merge(
        chars_lagged,
        on=[entity_col, date_col],
        how="inner",
    )

    return merged
