"""Chen-Pelger-Zhu Academic Asset Pricing Dataset — explore anonymized firm characteristics for ML benchmarking."""

import warnings

warnings.filterwarnings("ignore")

from datetime import datetime

import numpy as np
import plotly.express as px
import plotly.graph_objects as go
import polars as pl
import statsmodels.api as sm

from data import load_firm_characteristics
from utils.reproducibility import set_global_seeds
from utils.style import COLORS, ml4t_diverging

# Importing utils.style registers and activates the ML4T Plotly template
# (palette, gridlines, fonts) repo-wide; px/go figures below inherit its colorway.

# Production defaults — Papermill injects overrides for CI
SEED = 42

set_global_seeds(SEED)

# Load full dataset with split labels
df = load_firm_characteristics(split="all")

print(f"Total observations: {len(df):,}")
print(f"Columns: {len(df.columns)}")
df.columns

# Check date range and split distribution
date_stats = df.group_by("split").agg(
    pl.col("timestamp").min().alias("start"),
    pl.col("timestamp").max().alias("end"),
    pl.len().alias("n_obs"),
)
date_stats.sort("start")

# Observations per month
monthly_counts = (
    df.group_by("timestamp")
    .len()
    .sort("timestamp")
    .with_columns(
        pl.col("timestamp").dt.year().alias("year"),
    )
)

print("\nObservations per month:")
print(f"  Min: {monthly_counts['len'].min()}")
print(f"  Max: {monthly_counts['len'].max()}")
print(f"  Mean: {monthly_counts['len'].mean():.0f}")

# Plot observations over time
fig = px.line(
    monthly_counts.to_pandas(),
    x="timestamp",
    y="len",
    title="Coverage peaked at ~2,800 stocks in 2005, then eased to ~1,900 by 2016",
    labels={"len": "Stocks in cross-section", "timestamp": "Date"},
    color_discrete_sequence=[COLORS["blue"]],
)
# Split boundaries (train | valid | test), drawn in the neutral palette color
fig.add_vline(x=datetime(1990, 1, 1), line_dash="dash", line_color=COLORS["neutral"])
fig.add_vline(x=datetime(2000, 1, 1), line_dash="dash", line_color=COLORS["neutral"])
fig.add_annotation(
    x=datetime(1990, 1, 1),
    y=2800,
    text="Valid",
    showarrow=False,
    font=dict(color=COLORS["neutral"]),
)
fig.add_annotation(
    x=datetime(2000, 1, 1), y=2800, text="Test", showarrow=False, font=dict(color=COLORS["neutral"])
)
fig.show()

# Define characteristic categories
CHARACTERISTIC_CATEGORIES = {
    "Valuation": ["BEME", "E2P", "S2P", "CF2P", "D2P", "A2ME", "Q"],
    "Profitability": ["ROA", "ROE", "PROF", "OP", "PM", "PCM", "NI", "RNA"],
    "Investment": ["Investment", "NOA", "OA", "AC", "AT", "D2A"],
    "Momentum": ["r2_1", "r12_2", "r12_7", "r36_13", "ST_REV", "LT_Rev", "Rel2High"],
    "Risk": ["Beta", "MktBeta", "IdioVol", "Variance", "Resid_Var"],
    "Liquidity/Size": ["LME", "LTurnover", "Spread", "SUV"],
    "Leverage": ["Lev", "OL", "FC2Y", "C", "CF", "DPI2A"],
    "Other": ["ATO", "CTO", "SGA2S"],
}

# Print category summary
print("Characteristic Categories:")
print("=" * 60)
total = 0
for category, features in CHARACTERISTIC_CATEGORIES.items():
    print(f"\n{category} ({len(features)} features):")
    print(f"  {', '.join(features)}")
    total += len(features)
print(f"\nTotal: {total} characteristics")

# Distribution of each characteristic (confirming rank normalization)
feature_cols = [c for c in df.columns if c not in ["timestamp", "ret", "split"]]

# Compute statistics efficiently using Polars - no loops needed
sample_features = feature_cols[:10]
stats_df = df.select(sample_features).describe()
stats_df.filter(pl.col("statistic").is_in(["min", "max", "mean", "std"]))

# Compute correlation matrix
corr_data = df.select(feature_cols).to_numpy()
corr_matrix = np.corrcoef(corr_data, rowvar=False)

# Find highly correlated pairs
high_corr_pairs = []
for i in range(len(feature_cols)):
    for j in range(i + 1, len(feature_cols)):
        if abs(corr_matrix[i, j]) > 0.5:
            high_corr_pairs.append(
                {
                    "feature_1": feature_cols[i],
                    "feature_2": feature_cols[j],
                    "correlation": corr_matrix[i, j],
                }
            )

high_corr_df = pl.DataFrame(high_corr_pairs).sort("correlation", descending=True)
high_corr_df

# Correlation heatmap on the ML4T diverging scale (negative -> neutral -> positive)
_div = ml4t_diverging()
fig = go.Figure(
    data=go.Heatmap(
        z=corr_matrix,
        x=feature_cols,
        y=feature_cols,
        colorscale=[[0.0, _div[0]], [0.5, _div[1]], [1.0, _div[2]]],
        zmid=0,
        zmin=-1,
        zmax=1,
    )
)
fig.update_layout(
    title="Volatility and valuation form the densest correlation clusters",
    width=1000,
    height=900,
    xaxis_tickangle=-45,
)
fig.show()

# Step 1: Compute IC (correlation with returns) for each characteristic, per month.
# Sort by timestamp before dropping it: the Newey-West HAC t-stat below reads this
# as a time series, so the monthly ICs must be in chronological order. group_by emits
# groups in a non-deterministic order, which would scramble the autocovariance the HAC
# correction is built on (and make the t-stat non-reproducible run to run).
monthly_ics = (
    df.group_by("timestamp")
    .agg([pl.corr(col, "ret").alias(col) for col in feature_cols])
    .sort("timestamp")
    .drop("timestamp")
)

# Step 2: Compute mean IC, i.i.d. t-stat, and Newey-West HAC t-stat across months
n_months = len(monthly_ics)
# Newey-West lag selection: floor(4 * (T/100)^(2/9)); ~6 for monthly series of
# this length. Use 12 to be conservative against the annual cycle in factor returns.
nw_maxlags = 12
ic_stats = []
for col in feature_cols:
    ic_values = monthly_ics[col].to_numpy()
    ic_values = ic_values[~np.isnan(ic_values)]  # constant-feature months
    mean_ic = float(np.mean(ic_values))
    std_ic = float(np.std(ic_values, ddof=1))
    t_stat_iid = mean_ic / (std_ic / np.sqrt(len(ic_values)))
    # HAC t-stat: regress IC_t on a constant with Newey-West covariance.
    nw = sm.OLS(ic_values, np.ones(len(ic_values))).fit(
        cov_type="HAC", cov_kwds={"maxlags": nw_maxlags}
    )
    t_stat_nw = float(nw.tvalues[0])
    ic_stats.append(
        {
            "feature": col,
            "IC": mean_ic,
            "IC_std": std_ic,
            "t_stat_iid": t_stat_iid,
            "t_stat_NW": t_stat_nw,
            "abs_IC": abs(mean_ic),
        }
    )

ic_df = pl.DataFrame(ic_stats).sort("IC", descending=True)

print(
    f"Cross-sectional IC across {n_months} monthly cross-sections "
    f"(|t_NW| > 2 ≈ significant; Newey-West with {nw_maxlags} lags):"
)
ic_df

# Visualize ICs by category
ic_with_category = []
for row in ic_df.iter_rows(named=True):
    category = "Other"
    for cat, features in CHARACTERISTIC_CATEGORIES.items():
        if row["feature"] in features:
            category = cat
            break
    ic_with_category.append({**row, "category": category})

ic_cat_df = pl.DataFrame(ic_with_category).with_columns(
    pl.when(pl.col("IC") > 0)
    .then(pl.lit("Positive predictor"))
    .otherwise(pl.lit("Negative predictor"))
    .alias("Direction")
)

# 8 characteristic categories exceed the ML4T palette's distinct-color budget, so
# color encodes the sign of the relationship (the point of the figure); the x-axis
# names the individual characteristic.
fig = px.bar(
    ic_cat_df.to_pandas(),
    x="feature",
    y="IC",
    color="Direction",
    color_discrete_map={
        "Positive predictor": COLORS["blue"],
        "Negative predictor": COLORS["copper"],
    },
    category_orders={"Direction": ["Positive predictor", "Negative predictor"]},
    title="Short-term reversal and 12-month momentum carry the largest single-characteristic ICs",
    labels={"IC": "Information Coefficient", "feature": "Characteristic"},
)
fig.update_layout(xaxis_tickangle=-45, height=500)
fig.add_hline(y=0, line_dash="dash", line_color=COLORS["neutral"])
fig.show()

# Return statistics by split
return_stats = df.group_by("split").agg(
    pl.col("ret").mean().alias("mean"),
    pl.col("ret").std().alias("std"),
    pl.col("ret").min().alias("min"),
    pl.col("ret").max().alias("max"),
    pl.col("ret").quantile(0.25).alias("q25"),
    pl.col("ret").quantile(0.75).alias("q75"),
)
return_stats

# Return distribution
sample_size = min(100_000, len(df))
# Clip to +/-60% so the distribution shape is visible; a few extreme outliers
# (max next-month return ~+1900% in the valid split) otherwise flatten every bar
# into a single spike at zero and hide the very widening the figure is about.
fig = px.histogram(
    df.sample(n=sample_size, seed=SEED).to_pandas(),  # Convert for Plotly compatibility
    x="ret",
    color="split",
    nbins=120,
    range_x=[-0.6, 0.6],
    category_orders={"split": ["train", "valid", "test"]},
    color_discrete_sequence=[COLORS["blue"], COLORS["amber"], COLORS["copper"]],
    title="Return distributions widen after 1990 as the cross-section expands",
    labels={"ret": "Next-month return", "count": "Frequency"},
    opacity=0.6,
    barmode="overlay",
)
fig.show()

df_with_macro = load_firm_characteristics(split="all", include_macro=True)
macro_cols = [c for c in df_with_macro.columns if c.startswith("macro_")]
print(f"macro_* columns in shipped dataset: {len(macro_cols)}")