import os
import json
import gc
import pandas as pd
import numpy as np
from tqdm import tqdm
from numerapi import NumerAPI
from sklearn.model_selection import cross_val_score
from sklearn.linear_model import Ridge
import matplotlib.pyplot as plt
import seaborn as sns
from itertools import combinations

# Initialize NumerAPI
api = NumerAPI()

# Ensure the data directory exists
os.makedirs("../data", exist_ok=True)

# Download datasets
dataset_files = {
    "train.parquet": "v5.0/train.parquet",
    "validation.parquet": "v5.0/validation.parquet",
    "validation_example_preds.parquet": "v5.0/validation_example_preds.parquet",
    "features.json": "v5.0/features.json",
    "train_benchmark_models.parquet": "v5.0/train_benchmark_models.parquet",
    "validation_benchmark_models.parquet": "v5.0/validation_benchmark_models.parquet"
}

for filename, dataset in dataset_files.items():
    api.download_dataset(dataset, f"../data/{filename}")
    print(f"Downloaded {filename}")

# Load features metadata
with open("../data/features.json") as f:
    features_meta = json.load(f)

# Select the "all" feature set
features_to_load = features_meta["feature_sets"]["all"]

# Load datasets
train = pd.read_parquet(
    "../data/train.parquet",
    columns=features_to_load + ["era", "data_type", "target"]
)
validation = pd.read_parquet(
    "../data/validation.parquet",
    columns=features_to_load + ["era", "data_type", "target"]
)

# Load example predictions and meta models if needed
validation_example_preds = pd.read_parquet("../data/validation_example_preds.parquet")
train_benchmark_models = pd.read_parquet("../data/train_benchmark_models.parquet")
validation_benchmark_models = pd.read_parquet("../data/validation_benchmark_models.parquet")

# Process feature and target names
feature_names = [f for f in train.columns if "feature_" in f]
feature_names = [f for f in feature_names if f in features_to_load]
target_names = ["target"]

TARGET_NAME = target_names[0]
PREDICTION_NAME = "prediction"

# Drop rows with missing target values in validation set
validation = validation.dropna(subset=[TARGET_NAME], axis=0).copy()
gc.collect()

# Convert 'era' column to integer codes
train["era_int"] = train["era"].astype("category").cat.codes
validation["era_int"] = validation["era"].astype("category").cat.codes
gc.collect()

# Fill missing values in features with -2
for feature_name in tqdm(feature_names, desc="Filling missing feature values"):
    train[feature_name] = train[feature_name].fillna(-2)
    validation[feature_name] = validation[feature_name].fillna(-2)

# Fill missing values in targets with 0.5 (if applicable)
for target_name in tqdm(target_names, desc="Filling missing target values"):
    train[target_name] = train[target_name].fillna(0.5)
    validation[target_name] = validation[target_name].fillna(0.5)

gc.collect()

print(f"Feature set: {len(feature_names)} features")

# ----------------------------
# Generate Additional Features
# ----------------------------

# 1. Basic Statistical Features
def generate_basic_statistical_features(df_train, df_validation, feature_names):
    # Dictionaries to hold new columns
    train_stats = {}
    validation_stats = {}

    for feature in feature_names:
        train_stats[f'{feature}_mean'] = df_train[feature].mean()
        train_stats[f'{feature}_median'] = df_train[feature].median()
        train_stats[f'{feature}_std'] = df_train[feature].std()
        train_stats[f'{feature}_abs'] = df_train[feature].abs()

        # For validation, use training statistics
        validation_stats[f'{feature}_mean'] = df_train[feature].mean()
        validation_stats[f'{feature}_median'] = df_train[feature].median()
        validation_stats[f'{feature}_std'] = df_train[feature].std()
        validation_stats[f'{feature}abs'] = df_validation[feature].abs()

    # Convert dictionaries to DataFrames
    train_stats_df = pd.DataFrame(train_stats, index=df_train.index)
    validation_stats_df = pd.DataFrame(validation_stats, index=df_validation.index)

    # Concatenate to original DataFrames
    df_train = pd.concat([df_train, train_stats_df], axis=1)
    df_validation = pd.concat([df_validation, validation_stats_df], axis=1)

    return df_train, df_validation

print("Generating Basic Statistical Features...")
train, validation = generate_basic_statistical_features(train, validation, feature_names)
gc.collect()

# 2. Interaction Features
def generate_interaction_features(df_train, df_validation, feature_names):
    interaction_cols = []
    interaction_train = {}
    interaction_validation = {}

    total_combinations = len(feature_names) * (len(feature_names) - 1) // 2
    tqdm_combinations = tqdm(combinations(feature_names, 2), total=total_combinations, desc="Generating interaction features")

    for f1, f2 in tqdm_combinations:
        new_col_name = f'{f1}x{f2}'
        interaction_train[new_col_name] = df_train[f1] * df_train[f2]
        interaction_validation[new_col_name] = df_validation[f1] * df_validation[f2]

    # Convert dictionaries to DataFrames
    interaction_train_df = pd.DataFrame(interaction_train, index=df_train.index)
    interaction_validation_df = pd.DataFrame(interaction_validation, index=df_validation.index)

    # Concatenate to original DataFrames
    df_train = pd.concat([df_train, interaction_train_df], axis=1)
    df_validation = pd.concat([df_validation, interaction_validation_df], axis=1)

    return df_train, df_validation

print("Generating Interaction Features...")
train, validation = generate_interaction_features(train, validation, feature_names)
gc.collect()

# 3. Lag Features
def generate_lag_features(df_train, df_validation, feature_names, lags, group_col='era_int'):
    lag_train = {}
    lag_validation = {}

    for feature in feature_names:
        for lag in lags:
            lag_col_name = f'{feature}lag{lag}'
            lag_train[lag_col_name] = df_train.groupby(group_col)[feature].shift(lag)
            lag_validation[lag_col_name] = df_validation.groupby(group_col)[feature].shift(lag)

    # Convert dictionaries to DataFrames
    lag_train_df = pd.DataFrame(lag_train, index=df_train.index)
    lag_validation_df = pd.DataFrame(lag_validation, index=df_validation.index)

    # Concatenate to original DataFrames
    df_train = pd.concat([df_train, lag_train_df], axis=1)
    df_validation = pd.concat([df_validation, lag_validation_df], axis=1)

    # Drop rows with any NaN values resulted from lagging
    df_train = df_train.dropna().reset_index(drop=True)
    df_validation = df_validation.dropna().reset_index(drop=True)

    return df_train, df_validation

print("Generating Lag Features...")
lags = [1, 2, 3]
train, validation = generate_lag_features(train, validation, feature_names, lags)
gc.collect()

# Defragment DataFrames by copying
print("Defragmenting DataFrames...")
train = train.copy()
validation = validation.copy()
gc.collect()

# Update feature list
all_features = [col for col in train.columns if col not in ['era', 'data_type', 'target', 'era_int']]

# ----------------------------
# Feature Selection
# ----------------------------

# Define Performance Metric
def evaluate_performance(X, y, model):
    scores = cross_val_score(model, X, y, cv=5, scoring='r2', n_jobs=-1)
    return scores.mean()

# Feature Selection Algorithm
target = TARGET_NAME
candidate_features = all_features.copy()

selected_features = []
remaining_features = candidate_features.copy()
performance_history = []

model = Ridge(alpha=1.0)

y = train[target].values

print("Starting Feature Selection...")
while remaining_features:
    best_feature = None
    best_performance = -np.inf
    temp_performances = {}

    # Evaluate all remaining features in parallel
    tqdm_features = tqdm(remaining_features, desc="Evaluating features", leave=False)
    for feature in tqdm_features:
        current_features = selected_features + [feature]
        X = train[current_features].values
        performance = evaluate_performance(X, y, model)
        temp_performances[feature] = performance
        if performance > best_performance:
            best_performance = performance
            best_feature = feature

    if best_feature is None:
        break
    selected_features.append(best_feature)
    remaining_features.remove(best_feature)
    performance_history.append(best_performance)
    print(f"Selected Feature: {best_feature}, Current Performance: {best_performance:.4f}")

    # Check for improvement
    if len(performance_history) > 1:
        improvement = performance_history[-1] - performance_history[-2]
        if improvement < 0.001:
            print(f"Improvement {improvement:.6f} is less than threshold. Stopping feature selection.")
            break

# Evaluate and Select Features
feature_performance = pd.DataFrame({
    'Feature': selected_features,
    'Performance': performance_history
})

print("\nFeature Selection Performance:")
print(feature_performance)

# Plot Feature Selection Progress
plt.figure(figsize=(12, 6))
sns.lineplot(x=range(1, len(performance_history)+1), y=performance_history, marker='o')
plt.title('Feature Selection Progress')
plt.xlabel('Number of Features Selected')
plt.ylabel('Cross-Validated R²')
plt.grid(True)
plt.show()

# Analyzing the Results
final_performance = performance_history[-1] if performance_history else None
if final_performance is not None:
    print(f"\nFinal Performance with {len(selected_features)} Features: R² = {final_performance:.4f}")
else:
    print("\nNo features were selected.")

if selected_features:
    model.fit(train[selected_features], y)
    coefficients = pd.Series(model.coef, index=selected_features).sort_values(ascending=False)

    print("\nFeature Coefficients:")
    print(coefficients)

    # Plot Feature Coefficients
    plt.figure(figsize=(12, 6))
    coefficients.plot(kind='bar')
    plt.title('Feature Coefficients in Ridge Regression')
    plt.xlabel('Features')
    plt.ylabel('Coefficient Value')
    plt.tight_layout()
    plt.show()
else:
    print("No features were selected to fit the model.") This is my jupuyter notebook, optimize it for less memory usage, add checkpoints to resume work, because is consumption is too high (800gb ram)