"""CTF submission interface utilities.

Provides run_local() for executing a model script against local CTF data,
and generate_main_template() for wrapping arbitrary functions into the
exact CTF main() signature.
"""

from __future__ import annotations

import importlib.util
import sys
from collections.abc import Callable
from pathlib import Path
from typing import Any

import pandas as pd

# ---------------------------------------------------------------------------
# CTF pre-installed package list (Rule 8)
# ---------------------------------------------------------------------------

CTF_PREINSTALLED: set[str] = {
    "pandas",
    "numpy",
    "pyarrow",
    "boto3",
    "scipy",
    "scikit-learn",
    "sklearn",
    "polars",
    "joblib",
    # stdlib — never need to declare these
    "os", "sys", "time", "pathlib", "datetime", "traceback", "importlib",
    "math", "re", "json", "csv", "io", "abc", "collections", "itertools",
    "functools", "typing", "dataclasses", "enum", "copy", "warnings",
    "logging", "random", "hashlib", "struct", "gc", "threading",
    "multiprocessing", "concurrent", "contextlib", "string", "textwrap",
}


# ---------------------------------------------------------------------------
# Local execution
# ---------------------------------------------------------------------------


def run_local(
    script_path: str | Path,
    data_dir: str | Path,
    chars_file: str = "ctff_chars.parquet",
    features_file: str = "ctff_features.parquet",
    daily_ret_file: str = "ctff_daily_ret.parquet",
) -> pd.DataFrame:
    """Execute a CTF model script locally and return portfolio weights.

    Dynamically loads the script, reads the three CTF DataFrames from
    ``data_dir``, calls ``main(chars, features, daily_ret)``, and returns
    the resulting weights DataFrame.

    Parameters
    ----------
    script_path : str | Path
        Path to the Python model script (must define ``main()``).
    data_dir : str | Path
        Directory containing the three CTF Parquet files.
    chars_file, features_file, daily_ret_file : str
        File names within ``data_dir``.

    Returns
    -------
    pd.DataFrame
        DataFrame with columns ``id``, ``eom``, ``w``.

    Raises
    ------
    FileNotFoundError
        If script or any data file is missing.
    AttributeError
        If the script does not define a ``main`` function.
    """
    script_path = Path(script_path).resolve()
    data_dir = Path(data_dir).resolve()

    if not script_path.exists():
        raise FileNotFoundError(f"Script not found: {script_path}")

    for fname in (chars_file, features_file, daily_ret_file):
        p = data_dir / fname
        if not p.exists():
            raise FileNotFoundError(f"CTF data file not found: {p}")

    # Dynamically load the script as a module before reading data (fail fast).
    # Use path hash in name to avoid collisions between same-stem scripts.
    stem_hash = hex(hash(str(script_path)))[-6:]
    module_name = f"_ctf_model_{script_path.stem}_{stem_hash}"
    spec = importlib.util.spec_from_file_location(module_name, script_path)
    if spec is None or spec.loader is None:
        raise ImportError(f"Cannot load script: {script_path}")
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    try:
        spec.loader.exec_module(module)
    except Exception:
        sys.modules.pop(module_name, None)
        raise

    if not hasattr(module, "main"):
        sys.modules.pop(module_name, None)
        raise AttributeError(f"Script {script_path.name} does not define a main() function.")

    # Load data only after confirming the script is valid.
    chars = pd.read_parquet(data_dir / chars_file)
    features = pd.read_parquet(data_dir / features_file)
    daily_ret = pd.read_parquet(data_dir / daily_ret_file)

    weights = module.main(chars, features, daily_ret)

    # Normalise to pandas if Polars was returned
    if hasattr(weights, "to_pandas"):
        weights = weights.to_pandas()

    return weights  # type: ignore[no-any-return]


# ---------------------------------------------------------------------------
# Predictions → CTF weights converter
# ---------------------------------------------------------------------------

_VALID_MODES = frozenset({"long_short", "long_only", "zscore", "rank_weighted"})


def predictions_to_ctf_weights(
    predictions: pd.DataFrame,
    signal_col: str,
    id_col: str = "permno",
    date_col: str = "date",
    mode: str = "zscore",
    n_stocks: int = 100,
) -> pd.DataFrame:
    """Convert a predictions DataFrame to CTF portfolio weights.

    Accepts any flat DataFrame with a signal column and produces a
    CTF-compatible weights DataFrame with columns ``id``, ``eom``, ``w``.
    Use this when your predictions come from a source other than
    ``OOSResult`` (e.g. ``ParametricPolicy.weights()`` or a custom model).

    Parameters
    ----------
    predictions : pd.DataFrame
        Must contain ``id_col``, ``date_col``, and ``signal_col``.
    signal_col : str
        Column to use as the portfolio signal (predicted returns or scores).
    id_col : str
        Column name for the stock identifier.  Default ``"permno"``.
    date_col : str
        Column name for the formation date.  Default ``"date"``.
    mode : str
        Weighting scheme.  One of:

        ``"zscore"`` (default) — z-score weights normalized so
        ``sum(abs(w)) = 1``; mean-zero long-short over all stocks.

        ``"rank_weighted"`` — rank-centered weights normalized so
        ``sum(abs(w)) = 1``; mean-zero long-short over all stocks.

        ``"long_short"`` — equal-weight top/bottom ``n_stocks`` (weights
        sum to 0).

        ``"long_only"`` — equal-weight top ``n_stocks`` (weights sum to 1).
    n_stocks : int
        Number of stocks per leg for ``"long_short"`` / ``"long_only"``.
        Ignored for ``"zscore"`` and ``"rank_weighted"``.  Default 100.

    Returns
    -------
    pd.DataFrame
        Columns: ``id`` (int), ``eom`` (YYYY-MM-DD string), ``w`` (float).
        No duplicate (id, eom) pairs.
    """
    if mode not in _VALID_MODES:
        raise ValueError(f"mode must be one of {set(_VALID_MODES)}, got '{mode}'")

    if predictions.empty:
        return pd.DataFrame(columns=["id", "eom", "w"])

    records: list[dict[str, Any]] = []

    for date_val, group in predictions.groupby(date_col):
        eom_str = str(pd.Timestamp(date_val).date())  # type: ignore[arg-type]
        signal = group[signal_col]
        ids = group[id_col].astype(int).tolist()

        if mode == "zscore":
            mu = float(signal.mean())
            sigma = float(signal.std(ddof=1))
            if sigma == 0 or signal.isna().all():
                continue
            z = (signal - mu) / sigma
            abs_sum = float(z.abs().sum())
            if abs_sum == 0:
                continue
            for stock_id, w in zip(ids, (z / abs_sum).tolist()):
                records.append({"id": stock_id, "eom": eom_str, "w": w})

        elif mode == "rank_weighted":
            r = signal.rank(method="average", ascending=True)
            centered = r - float(r.mean())
            abs_sum = float(centered.abs().sum())
            if abs_sum == 0:
                continue
            for stock_id, w in zip(ids, (centered / abs_sum).tolist()):
                records.append({"id": stock_id, "eom": eom_str, "w": w})

        else:
            ranked = signal.rank(method="first", ascending=False)
            if mode == "long_short":
                top_mask = ranked <= n_stocks
                bot_mask = (ranked > (len(group) - n_stocks)) & (~top_mask)
                n_long = int(top_mask.sum())
                n_short = int(bot_mask.sum())
                if n_long == 0:
                    continue
                for idx in group.index[top_mask]:
                    records.append({"id": int(group.loc[idx, id_col]),
                                    "eom": eom_str, "w": 1.0 / n_long})
                for idx in group.index[bot_mask]:
                    records.append({"id": int(group.loc[idx, id_col]),
                                    "eom": eom_str, "w": -1.0 / n_short})
            else:  # long_only
                top_mask = ranked <= n_stocks
                n_long = int(top_mask.sum())
                if n_long == 0:
                    continue
                for idx in group.index[top_mask]:
                    records.append({"id": int(group.loc[idx, id_col]),
                                    "eom": eom_str, "w": 1.0 / n_long})

    result = pd.DataFrame(records, columns=["id", "eom", "w"])
    result["id"] = result["id"].astype(int)
    return result


# ---------------------------------------------------------------------------
# main() template generator
# ---------------------------------------------------------------------------


def generate_main_template(model_fn: Callable[..., Any]) -> str:
    """Generate a CTF-compliant main() wrapper around an arbitrary function.

    Returns a Python source string that defines ``main(chars, features,
    daily_ret)`` delegating to ``model_fn``.

    Parameters
    ----------
    model_fn : Callable
        The user-defined portfolio construction function.  Must accept
        (chars, features, daily_ret) or a subset.

    Returns
    -------
    str
        Python source code string ready to be written to a .py file.
        The caller is responsible for prepending any import statements
        needed to make ``model_fn`` available in the generated script.
    """
    fn_name = model_fn.__name__
    return (
        "import pandas as pd\n\n\n"
        f"def main(\n"
        f"    chars: pd.DataFrame,\n"
        f"    features: pd.DataFrame,\n"
        f"    daily_ret: pd.DataFrame,\n"
        f") -> pd.DataFrame:\n"
        f'    """CTF submission: delegates to {fn_name}."""\n'
        f"    return {fn_name}(chars, features, daily_ret)\n"
    )
